#include <ElfParser.h>
#include <MonitorDriver.h>
#include <Filesystem.h>
#include <FileDescriptor.h>
#include <Kernel.h>
#include <ProcessManager.h>

Elf32Parser::Elf32Parser()
{
  header             = NULL;
  sectionHeaders     = NULL;
  symbolTable        = NULL;
  stringTable        = NULL;
  gotTable           = NULL;
  filename           = NULL;
}

Elf32Parser::Elf32Parser(char *fname)
{
  header             = NULL;
  sectionHeaders     = NULL;
  symbolTable        = NULL;
  stringTable        = NULL;
  gotTable           = NULL;
  filename           = NULL;

  filename = new char[Kernel::strlen(fname)+1];
  Kernel::strcpy(filename, fname);

  // Firstly, make sure the file can be read. We do this via a static file,
  // so we don't use up one of the supercool stdin/stdout/stderr descriptors by
  // accident
  File *file = filesystem->traverse(filename);
  if (!file)
  {
    delete [] filename;
    filename = NULL;
    return;
  }
  file->open(NULL);
  
  // Load in the file header.
  header = new Elf32Header;
  
  u8int *buffer = (u8int*)memoryManager.umalloc(sizeof(Elf32Header));
  
  u32int result = file->read(0, sizeof(Elf32Header), buffer);
  memcpy((u8int*)header, buffer, sizeof(Elf32Header));
  
  memoryManager.ufree(buffer);
  
  sectionHeaders = new Elf32SectionHeader[header->shnum];

  ASSERT(header->shentsize == sizeof(Elf32SectionHeader));
  buffer = (u8int*)memoryManager.umalloc(sizeof(Elf32SectionHeader)*header->shnum);
  result = file->read(header->shoff,
                      header->shnum*header->shentsize,
                      buffer);
  memcpy((u8int*)sectionHeaders, buffer, sizeof(Elf32SectionHeader)*header->shnum);
  memoryManager.ufree(buffer);
  
  stringTable = &(sectionHeaders[header->shstrndx]);

  for (int i = 0; i < header->shnum; i++)
  {
    switch(sectionHeaders[i].type)
    {
      case SHT_SYMTAB:
        symbolTable = &(sectionHeaders[i]);
        break;
    }
  }
  
  // Load the string table.
  
  char *strtab = (char*)memoryManager.umalloc(stringTable->size);
  file->read(stringTable->offset, stringTable->size, (u8int*)strtab);
  for (int i = 0; i < header->shnum; i++)
  {
    char *str = (char *) strtab + sectionHeaders[i].name;
    if (!Kernel::strcmp(str, ".got.plt"))
    {
      gotTable = &(sectionHeaders[i]);
    }
    if (!Kernel::strcmp(str, ".rel.plt"))
    {
      relTable = &(sectionHeaders[i]);
    }
  }
  memoryManager.ufree(strtab);

}

Elf32Parser::~Elf32Parser()
{
  delete header;
  delete sectionHeaders;
  delete [] filename;
}

void Elf32Parser::writeAllSections()
{
  File *file = filesystem->traverse(filename);
  ASSERT(file);
  file->open(NULL);

  for (int i = 0; i < header->shnum; i++)
  {
    if (sectionHeaders[i].flags & SHF_ALLOC)
    {
      memoryManager.allocateRange(sectionHeaders[i].addr,
                                  sectionHeaders[i].size);
      if (sectionHeaders[i].type != SHT_NOBITS)
      {
        // Copy section data from the file.
        u8int *buf = (u8int*)memoryManager.umalloc(sectionHeaders[i].size);
        u32int ret = file->read(sectionHeaders[i].offset,
                                sectionHeaders[i].size,
                                buf);
        DEBUG_MSG("*** Writing at " << hex << (u32int)sectionHeaders[i].addr);
//         if (strcmp(filename, "/mnt/bin/bash"))
        Kernel::memcpy((u8int*)sectionHeaders[i].addr,
                       buf,
                       sectionHeaders[i].size);
        memoryManager.ufree(buf);
      }
      else
      {
        DEBUG_MSG("*** Memsetting at " << hex << (u32int)sectionHeaders[i].addr);
//         if (strcmp(filename, "/mnt/bin/bash"))
        Kernel::memset((u8int*)sectionHeaders[i].addr,
                       0,
                       sectionHeaders[i].size);
      }
    }
  } 
}

void Elf32Parser::loadKernel(Elf32SectionHeader *symtab, Elf32SectionHeader *strtab)
{
  symbolTable = symtab;
  stringTable = strtab;
}

char *Elf32Parser::findSymbol(Address addr, Address *symbolStart)
{
  Elf32Symbol *symbol = (Elf32Symbol *)symbolTable->addr;
  
  int i;
  for(i = 0; i < symbolTable->size / sizeof(Elf32Symbol); i++)
  {
    if ( (addr >= symbol->value) &&
         (addr <  symbol->value+symbol->size) )
    {
      char *c = (char *)(symbol->name)+stringTable->addr;
     
      if (symbolStart)
      {
        *symbolStart = symbol->value;
      }
      return c;
    }
    symbol ++;
  }
  return NULL;

}

Address Elf32Parser::findSymbol(char *str)
{
  Elf32Symbol *symbol = (Elf32Symbol *)symbolTable->addr;
  
  int i;
  for (i = 0; i < symbolTable->size / sizeof(Elf32Symbol); i++)
  {
    char *c = (char *)(symbol->name)+stringTable->addr;
    
    if (!Kernel::strcmp(c, str))
    {
      return symbol->value;
    }
    symbol ++;
  }
  return 0;
}

Address Elf32Parser::findDynamicSymbolLocation(Address o)
{
  // For every section that is of type REL...
    Elf32SectionHeader *head = relTable;
    if (head->type == SHT_REL)
    {
      // find the correct rel entry.
      Elf32Rel *rel = (Elf32Rel *)(head->addr+o);
      return rel->offset;
    }
  return NULL;
}

char *Elf32Parser::findDynamicSymbolName(Address o)
{
  // For every section that is of type REL...
    Elf32SectionHeader *head = relTable;
    if (head->type == SHT_REL)
    {
      // find the correct rel entry.
      Elf32Rel *rel = (Elf32Rel *)(head->addr+o);
      // find the section index for the dynamic symbol table
      Elf32SectionHeader dyn = sectionHeaders[head->link];
      // find the dynamic tag for this symbol
      Elf32Dyn *tag = (Elf32Dyn*)(dyn.addr+ ELF32_R_SYM(rel->info)*dyn.entsize);
      // find the string table for this symbol
      Elf32SectionHeader strtab = sectionHeaders[dyn.link];
      return (char *)(strtab.addr + tag->tag);
    }
  return NULL;
}

Address Elf32Parser::getGlobalOffsetTable()
{
  if (gotTable)
  {
    return (Address)gotTable->addr;
  }
  else
  {
    return 0;
  }
}

Elf32Parser *Elf32Parser::clone()
{
  if (!filename)
  {
    return new Elf32Parser();
  }
  Elf32Parser *other = new Elf32Parser(filename);
  return other;
}

Address Elf32Parser::getLastAddress()
{
  // Iterate over all the sections, getting the maximum last address.
  Address max = 0;
  for (int i = 0; i < header->shnum; i++)
  {
    if (sectionHeaders[i].flags & SHF_ALLOC)
    {
      Address addr = sectionHeaders[i].addr + sectionHeaders[i].size;
      if (addr > max)
      {
        max = addr;
      }
    }
  }
  return max;
}
